# coding: utf-8
import numpy as np
from .geo_classes import Node, Region
import logging
from sklearn.svm import SVR
from sklearn.cross_validation import train_test_split
from sklearn import linear_model
import matplotlib.pyplot as plt
import pickle

try:
    from sklearn.cross_validation import KFold
    __USE_OLD_KOLD = True
except Exception:
    from sklearn.model_selection import KFold
    __USE_OLD_KOLD = False


def get_train_test(X, y, kfold=False, kfold_k=None, kfold_shuffle=False):
    if kfold is False:
        yield X, y, None, None
    else:
        if __USE_OLD_KOLD:
            kf = KFold(X.shape[0], n_folds=kfold_k, shuffle=kfold_shuffle)
            for train_index, test_index in kf:
                yield X[train_index], y[train_index], X[test_index], y[test_index]
        else:
            kf = KFold(n_splits=kfold_k, shuffle=kfold_shuffle)
            for train_index, test_index in kf.split(X):
                yield X[train_index], y[train_index], X[test_index], y[test_index]


class Model(object):
    def __init__(self, x_file, y_file):
        self.X = []
        self.y = []
        x_regions = []
        with open(x_file) as f:
            for line in f:
                line = line.split(' ')
                x_regions.append(Region(Node(float(line[0]), float(line[1])), Node(float(line[2]), float(line[3]))))
                self.X.append([float(i) for i in line[4:len(line)]])

        """x可能有缺，而y无缺，做一个对齐"""
        pointer = 0
        with open(y_file) as f:
            for line in f:
                line = line.split(' ')
                region = Region(Node(float(line[0]), float(line[1])), Node(float(line[2]), float(line[3])))
                if region == x_regions[pointer]:
                    self.y.append(float(line[4]))
                    pointer += 1

        self.X = np.array(self.X)
        self.y = np.array(self.y)
        logging.info(str(self.X.shape))
        logging.info(str(self.y.shape))

    def do_svr(self, train_factor=0.7, random_state=0):
        X_train, X_test, y_train, y_test = train_test_split(
            self.X, self.y, test_size=1 - train_factor, random_state=random_state)
        print(X_train.shape, X_test.shape, y_train.shape, y_test.shape)
        clf = SVR(C=10, epsilon=0.1)
        clf.fit(X_train, y_train)
        print(clf.score(X_test, y_test))
        y_predicted = clf.predict(X_test)
        fig, ax = plt.subplots()
        ax.scatter(y_test, y_predicted)
        ax.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'k--', lw=4)
        ax.set_xlabel('Measured')
        ax.set_ylabel('Predicted')
        plt.show()
        self.save_svr_model()

    def do_svr_kfold(self, kfold_k=5, vis=True):
        all_score = 0
        for i, (X_train, y_train, X_test, y_test) in enumerate(get_train_test(self.X, self.y, True, kfold_k, True)):
            print(X_train.shape, X_test.shape, y_train.shape, y_test.shape)
            clf = SVR(C=10, epsilon=0.1)
            clf.fit(X_train, y_train)
            score = clf.score(X_test, y_test)
            print('Score: %d / %d: %f' % (i + 1, kfold_k, score))
            all_score += score
            if vis:
                y_predicted = clf.predict(X_test)
                fig, ax = plt.subplots()
                ax.scatter(y_test, y_predicted)
                ax.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'k--', lw=4)
                ax.set_xlabel('Measured')
                ax.set_ylabel('Predicted')
                plt.show()
        print('avg score: %f' % (all_score / kfold_k))

    def save_svr_model(self):
        clf = SVR(C=10, epsilon=0.1)
        clf.fit(self.X, self.y)
        with open('data/svr.model', 'wb') as f:
            pickle.dump(clf, f)
        print('Model saved to: ' + 'data/svr.model')

    def do_linear_regression(self, train_factor=0.9, random_state=0):
        X_train, X_test, y_train, y_test = train_test_split(
            self.X, self.y, test_size=1 - train_factor, random_state=random_state)
        print(X_train.shape, X_test.shape, y_train.shape, y_test.shape)
        clf = linear_model.LinearRegression()
        clf.fit(X_train, y_train)
        print(clf.score(X_test, y_test))
